from implement.layers.basic.model import Model
from implement.layers.basic.batchnorm import BatchNorm
from implement.layers.conv.building_block import BuildingBlock
from implement.layers.conv.conv2d import Conv2d
from implement.layers.basic.affine import Affine
from utils.file_opera import get_file
from utils.functions_collect import average_pooling, reshape, relu, pooling

# 该模型是 ResNet，用于图像分类任务。可以选择加载预训练的权重。
class ResNet(Model):
    WEIGHTS_PATH = 'https://github.com/koki0702/dezero-models/releases/download/v0.1/resnet{}.npz'

    def __init__(self, n_layers=152, pretrained=False):
        """ResNet 模型的构造函数。

        Args:
            n_layers (int): ResNet 的层数，可选值为 50、101 或 152。
            pretrained (bool): 是否加载预训练的权重。
        """
        super().__init__()

        if n_layers == 50:
            block = [3, 4, 6, 3]
        elif n_layers == 101:
            block = [3, 4, 23, 3]
        elif n_layers == 152:
            block = [3, 8, 36, 3]
        else:
            raise ValueError('The n_layers argument should be either 50, 101,'
                             ' or 152, but {} was given.'.format(n_layers))

        self.conv1 = Conv2d(64, 7, 2, 3)
        self.bn1 = BatchNorm()
        self.res2 = BuildingBlock(block[0], 64, 64, 256, 1)
        self.res3 = BuildingBlock(block[1], 256, 128, 512, 2)
        self.res4 = BuildingBlock(block[2], 512, 256, 1024, 2)
        self.res5 = BuildingBlock(block[3], 1024, 512, 2048, 2)
        self.fc6 = Affine(1000)

        if pretrained:
            weights_path = get_file(ResNet.WEIGHTS_PATH.format(n_layers))
            self.load_weights(weights_path)

    def forward(self, x):
        """ResNet 模型的前向传播。

        Args:
            x (ndarray): 输入图像数据。

        Returns:
            ndarray: ResNet 模型的输出。
        """
        x = relu(self.bn1(self.conv1(x)))
        x = pooling(x, kernel_size=3, stride=2)
        x = self.res2(x)
        x = self.res3(x)
        x = self.res4(x)
        x = self.res5(x)
        x = _global_average_pooling_2d(x)
        x = self.fc6(x)
        return x

def _global_average_pooling_2d(x):
    """全局平均池化操作。

    Args:
        x (ndarray): 输入数据。

    Returns:
        ndarray: 池化后的结果。
    """
    N, C, H, W = x.shape
    h = average_pooling(x, (H, W), stride=1)
    h = reshape(h, (N, C))
    return h

